# coding:utf-8
import re
import json
import time
import redis
import numpy as np
from collections import Counter
from utils.tool import FpTool
from utils.opensearch_client import OpensearchClient
from utils.search_client import SearchClient
from utils.setting import OPENSEARCH_FIELDS
import datetime
import random

class SimilaritySearch(FpTool):
    '''排重 召回模块'''
    def __init__(self, index_name="group_id_online", logger="",is_debug=False):
        super(SimilaritySearch, self).__init__(logger, write_log_to_kafka = not is_debug)
        self.is_debug = is_debug
        if self.is_debug:
            self.debug_index = 'test_group_id'
            index_name = self.debug_index
        self.search_client = OpensearchClient(index=index_name)
        self.search_client_new = SearchClient(index=index_name)
        self.text_sim_docs_num = 50
        self.vec_sim_docs_num = 50
        # 必须字段
        self.fields={
            "title",
            "title_split",
            "title_entity",
            "keywords",
            "classes",
            "group_id",
            "ctime",
            "dataid"
        }

    @FpTool.cost
    def search(self):
        body = {"size":10000, "query": {"bool": {"must": [{"range": {"ctime": {"gte": 1650384000}}}]}}}
        for docs in self.search_client.large(body):
            for doc in docs:
                _ndata = self.__NorData(doc['_source'], OPENSEARCH_FIELDS)
                #print(_ndata)
                self.insert(doc['_source'])
                #exit()
            
    @FpTool.cost
    def insert(self, data):
        try:
            _ndata = self.__NorData(data, OPENSEARCH_FIELDS)
            # 插入 opensearch 库
            #st = time.time()
            ret = "default"
            if self.is_debug:
                ret = self.search_client_new.insert(_ndata, id=_ndata['dataid'], index_name=self.debug_index)
            else:
                try:
                    ctime = int(data.get('ctime', time.time()))
                except:
                    ctime = int(time.time())
                if abs(time.time() - ctime) < 3600*24*40:
                    format_ctime = time.localtime(ctime)
                    index_name = "group_id_{}_{:0>2d}".format(format_ctime.tm_year, format_ctime.tm_mon) # 按月建索引
                    #ret = self.search_client.insert(_ndata, id=_ndata['dataid'], index_name=index_name)
                    ret = self.search_client_new.insert(_ndata, id=_ndata['dataid'], index_name=index_name) ## 双写
                else: # 超过半年的文章不入索引
                    self.logger("old news skip insert opensearch! data:{}".format(data))
            #print("insert costs {:.4f}".format(1000*(time.time()-st)))
            self.logger("insert data to elasticsearch:{} data:{}".format(ret, _ndata))
        except Exception as e:
            self.logger("insert to db error! msg:{}".format(e))
            return False
        return True

    def __NorData(self, data, fields):
        return {key:data[key] for key in fields if key in data}


search_client = SimilaritySearch()
search_client.search()
